# -*- coding: utf-8 -*-
import time
import math
import datetime
from pymongo import MongoClient
from conf.settings import CLOUDMONITOR_MONGODB, CLOUDMONITOR_DB
from utility.log import Logger
from utility.wx_notify import notify
logger = Logger(filename='metrics.log').get_logger()

INTERVAL = 7
db = MongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]
MERGE = {'$merge': {'into': 'metrics', 'on': '_id', 'whenMatched': "replace", 'whenNotMatched': "insert"}}


def task_type_switch(key):
    return {
        '$switch': {
            'branches': [
                {'case': {'$eq': [key, 'vul']}, 'then': '漏洞监测'},
                {'case': {'$eq': [key, 'ssl']}, 'then': 'SSL监测'},
                {'case': {'$eq': [key, 'asset']}, 'then': '资产变动'},
                {'case': {'$eq': [key, 'content']}, 'then': '内容监测'},
                {'case': {'$eq': [key, 'http']}, 'then': 'HTTP监测'},
                {'case': {'$eq': [key, 'ping']}, 'then': 'PING监测'},
                {'case': {'$eq': [key, 'securityEvent']}, 'then': '安全事件'},
                {'case': {'$eq': [key, 'sitePortrait']}, 'then': '网站画像'},
                {'case': {'$eq': [key, 'ipv6']}, 'then': 'IPv6合规'},
                {'case': {'$eq': [key, 'hostVul']}, 'then': '主机监测'},
            ],
            'default': key
        }
    }


def level_switch(key):
    return {
        '$switch': {
            'branches': [
                {'case': {'$eq': [key, 1]}, 'then': '低危'},
                {'case': {'$eq': [key, 2]}, 'then': '中危'},
                {'case': {'$eq': [key, 3]}, 'then': '高危'}
            ],
            'default': key
        }
    }


def date_to_string(key):
    return {'$dateToString': {'date': key, 'format': '%Y%m%d'}}


def vul_level(start, end):
    match = {
        'endTime': {'$gt': start, '$lt': end},
        'taskType': {'$in': ['vul', 'ssl', 'content', 'securityEvent']}
    }
    group = {
        '_id': {
            'taskType': '$taskType',
            'level': {
                '$switch': {
                    'branches': [
                        {'case': {'$in': ['$taskType', ['vul', 'ssl']]}, 'then': '$result.warnings.level'}
                    ],
                    'default': '$result.warnings.title',
                }
            },
            'endTime': date_to_string('$endTime'),
            'taskTypeCn': task_type_switch('$taskType'),
            'metrics_type': 'risk_level'
        },
        'times': {'$sum': 1},
        'domain': {'$addToSet': '$domain'}
    }
    project = {
        'taskTypeCn': '$_id.taskTypeCn',
        'levelCn': level_switch('$_id.level'),
        'date': {'$toDate': '$_id.endTime'},
        'times': 1,
        'metrics_type': 'risk_level'
    }
    pipeline = [
        {'$match': match},
        {'$unwind': {'path': '$result.warnings'}},
        {'$group': group},
        {'$project': project},
        MERGE
    ]
    db.tasks.aggregate(pipeline)


def task_type_sum(start, end):
    match = {
        'endTime': {'$gt': start, '$lt': end},
    }
    group = {
        '_id': {
            'taskType': '$taskType',
            'endTime': date_to_string('$endTime'),
            'taskTypeCn': task_type_switch('$taskType'),
            'metrics_type': 'task_type_sum',
        },
        'times': {'$sum': 1}
    }
    project = {
        'times': 1,
        'taskTypeCn': '$_id.taskTypeCn',
        'metrics_type': 'task_type_sum',
        'date': {'$toDate': '$_id.endTime'},
    }
    pipeline = [
        {'$match': match},
        {'$group': group},
        {'$project': project},
        MERGE
    ]
    db.tasks.aggregate(pipeline)


def wx(start_date, start_time):
    ret = db.metrics.aggregate([
        {'$match': {'metrics_type': 'task_type_sum', 'date': {'$gte': start_date}}},
        {'$group': {'_id': '$taskTypeCn', 'sum': {'$sum': '$times'}}}])
    message = f"汇总数据结束, 用时{int(time.time()-start_time)}秒，" \
              f"汇总数据范围{start_date} - {datetime.datetime.now()}，"
    message += " ".join([f"{_['_id']}:{_['sum']}次" for _ in ret])
    logger.info(message)
    notify(message)


def main(start, end):
    logger.info("开始汇总统计数据......")
    start_time = time.time()
    END = end or datetime.datetime.now()
    START = start or datetime.datetime.strptime(datetime.datetime.utcnow().strftime('%Y-%m-%d'), '%Y-%m-%d') - datetime.timedelta(days=INTERVAL-1)
    days = (END - START).days
    for _ in range(math.ceil(days/INTERVAL)):
        start = START+datetime.timedelta(days=(_ * INTERVAL))
        end = min([END, START++datetime.timedelta(days=((_ + 1) * INTERVAL))])
        logger.info(f"开始汇总{start} ~ {end}")
        task_type_sum(start, end)
        vul_level(start, end)
    wx(START, start_time)
    logger.info("汇总统计数据完成......")


if __name__ == "__main__":
    # start = datetime.datetime(2022, 1, 1)
    # end = datetime.datetime(2022, 7, 1)
    start = end = None
    main(start, end)
